/*
 * File:          my_controller.c
 * Date:
 * Description:
 * Author:
 * Modifications:
 */

/*
 * You may need to add include files like <webots/distance_sensor.h> or
 * <webots/motor.h>, etc.
 */
#include <webots/robot.h>
#include <webots/motor.h>
#include <webots/keyboard.h>
#include <webots/inertial_unit.h>
#include <stdio.h>
#include <math.h>

/*
 * You may want to add macros here.dw
 */
#define TIME_STEP 64
#define MAX_SPEED 10.0  // 最大速度常量
#define MAX_STEER_ANGLE M_PI/4  // 最大转向角度45度


int main(int argc, char **argv) {
  /* necessary to initialize webots stuff */
  wb_robot_init();

  // 初始化键盘监听
  wb_keyboard_enable(TIME_STEP);
  
  
 // 获取电机设备
  WbDeviceTag drives[4] = {
      wb_robot_get_device("lf_drive"),
      wb_robot_get_device("rf_drive"),
      wb_robot_get_device("lb_drive"),
      wb_robot_get_device("rb_drive")
  };
  
  WbDeviceTag steers[4] = {
      wb_robot_get_device("lf_steer"),
      wb_robot_get_device("rf_steer"),
      wb_robot_get_device("lb_steer"),
      wb_robot_get_device("rb_steer")
  };
  
  
       // 配置驱动轮：速度控制模式
  for (int i = 0; i < 4; i++) {
      wb_motor_set_position(drives[i], INFINITY);  // 关键修复：启用速度控制
      wb_motor_set_velocity(drives[i], 0.0);       // 初始速度为0
  }

  // 配置转向轮：位置控制模式
  for (int i = 0; i < 4; i++) {
      wb_motor_set_position(steers[i], INFINITY);  // 启用位置控制
      wb_motor_set_position(steers[i], 0.0);       // 设置初始角度为0（直行）
  }
  
   // 运动状态变量
  double speed = 0.0;
  double steer_angle = 0.0;
  double rotate_speed = 0.0; // 新增旋转速度变量
  int strafe_mode = 0; // 0=正常模式, 1=左移, 2=右移
  double move_speed = 0.0; // 平移速度
  double move_dir_angle = 0.0; // 平移方向角（弧度，0为正前，PI/2为正右）
  // 记录自转状态
  int rotate_mode = 0; // 0=不自转, 1=自转
  double robot_theta = 0.0; // 机器人朝向（弧度）
  
  // 显示控制说明
  printf("\n=== Steering Wheel Chassis Control ===\n");
  printf("  W: 前进\n");
  printf("  S: 后退\n");
  printf("  A: 左平移\n");
  printf("  D: 右平移\n");
  printf("   可组合移动\n");
  printf("  X: 逆时针旋转\n");
  printf(" 小陀螺：按住x与其他组合\n");
  
  // 获取IMU设备
  WbDeviceTag imu = wb_robot_get_device("imu");
  wb_inertial_unit_enable(imu, TIME_STEP);
  
  while (wb_robot_step(TIME_STEP) != -1) {
    // 每帧先重置运动状态，松手即停
    move_speed = 0.0;
    strafe_mode = 0;
    rotate_mode = 0;

    // 记录多键状态
    int w_down = 0, a_down = 0, s_down = 0, d_down = 0, x_down = 0;
    int key;
    while ((key = wb_keyboard_get_key()) != -1) {
      if (key == 'W') w_down = 1;
      if (key == 'A') a_down = 1;
      if (key == 'S') s_down = 1;
      if (key == 'D') d_down = 1;
      if (key == 'X') x_down = 1;
    }
    // X优先级最高，按下X时自转+世界系WASD平移
    if (x_down) {
      // 读取IMU航向角
      const double *rpy = wb_inertial_unit_get_roll_pitch_yaw(imu);
      double yaw = rpy[2]; // 世界坐标系下的车体朝向
      // 方向合成与归一化
      double vx = 0.0, vy = 0.0;
      if (w_down) vy += 1.0; // 世界系Y正方向
      if (s_down) vy -= 1.0;
      if (a_down) vx -= 1.0;
      if (d_down) vx += 1.0;
      // 冲突检测
      if ((w_down && s_down) || (a_down && d_down)) {
        vx = 0.0; vy = 0.0;
      }
      double norm = sqrt(vx*vx + vy*vy);
      if (norm > 0.0) {
        vx /= norm;
        vy /= norm;
      }
      // 世界系速度旋转到车体系
      double vx_body =  cos(-yaw) * vx - sin(-yaw) * vy;
      double vy_body =  sin(-yaw) * vx + cos(-yaw) * vy;
      
      
      double move_speed =  MAX_SPEED;//平移速度
      double w = -60.0; // 自转角速度，负号实现逆时针
      
      
      double wheel_x[4] = {  0.2,  0.2, -0.2, -0.2 };
      double wheel_y[4] = {  0.2, -0.2,  0.2, -0.2 };
      double v_array[4];
      double v_max = 0.0;
      static double last_steer_angle[4] = {0, 0, 0, 0};
      for (int i = 0; i < 4; i++) {
        double vix = vx_body * move_speed - w * wheel_y[i];
        double viy = vy_body * move_speed + w * wheel_x[i];
        double angle = atan2(viy, vix);
        double v = sqrt(vix*vix + viy*viy);
        double current_angle = last_steer_angle[i];
        // 归一化到[-pi, pi]
        angle = fmod(angle + M_PI, 2*M_PI); if (angle < 0) angle += 2*M_PI; angle -= M_PI;
        current_angle = fmod(current_angle + M_PI, 2*M_PI); if (current_angle < 0) current_angle += 2*M_PI; current_angle -= M_PI;
        double delta = angle - current_angle;
        // 归一化delta到[-pi, pi]
        delta = fmod(delta + M_PI, 2*M_PI); if (delta < 0) delta += 2*M_PI; delta -= M_PI;
        double delta_flip = delta > 0 ? delta - M_PI : delta + M_PI;
        if (fabs(delta_flip) < fabs(delta) - 1e-3) {
          angle += (delta > 0 ? -M_PI : M_PI);
          v = -v;
          angle = fmod(angle + M_PI, 2*M_PI); if (angle < 0) angle += 2*M_PI; angle -= M_PI;
        }
        if (angle > M_PI/2) angle -= M_PI;
        if (angle < -M_PI/2) angle += M_PI;
        last_steer_angle[i] = angle;
        v_array[i] = v;
        if (fabs(v) > v_max) v_max = fabs(v);
        if (steers[i]) wb_motor_set_position(steers[i], angle);
      }
      double scale = 1.0;
      if (v_max > MAX_SPEED) scale = MAX_SPEED / v_max;
      for (int i = 0; i < 4; i++) {
        if (drives[i]) wb_motor_set_velocity(drives[i], v_array[i] * scale);
      }
      continue; // 跳过本帧后续处理
    }
    // 组合键方向逻辑
    move_speed = 0.0;
    strafe_mode = 0;
    double vx = 0.0, vy = 0.0;
    // 修正方向映射：W前进，S后退，A左移，D右移
    if (w_down) vy -= 1.0;    // W 前进（负Y）
    if (s_down) vy += 1.0;    // S 后退（正Y）
    if (a_down) vx -= 1.0;    // A 左移
    if (d_down) vx += 1.0;    // D 右移
    // 方向合成与归一化，保证斜向切换平滑
    vx = 0.0; vy = 0.0;
    if (w_down) vy -= 1.0;
    if (s_down) vy += 1.0;
    if (a_down) vx -= 1.0;
    if (d_down) vx += 1.0;
    // 冲突检测：W+S或A+D等，直接停止
    if ((w_down && s_down) || (a_down && d_down)) {
      vx = 0.0; vy = 0.0;
    }
    double norm = sqrt(vx*vx + vy*vy);
    if (norm > 0.0) {
      vx /= norm;
      vy /= norm;
      move_speed = MAX_SPEED;
      move_dir_angle = atan2(vx, vy);
      strafe_mode = 5;
    } else {
      move_speed = 0.0;
      strafe_mode = 0;
      // 不再自动回正舵轮，松开按键时保持当前舵轮角度和速度
      // 仅驱动轮速度归零
      for (int i = 0; i < 4; i++) {
        if (drives[i]) wb_motor_set_velocity(drives[i], 0.0);
      }
      continue; // 跳过本帧后续处理
    }
    // 计算机器人朝向（模拟积分，假设每步旋转速度恒定）
    if (rotate_mode) {
      robot_theta += 5.0 * TIME_STEP / 1000.0; // 5rad/s, TIME_STEP单位ms
      if (robot_theta > M_PI) robot_theta -= 2*M_PI;
      if (robot_theta < -M_PI) robot_theta += 2*M_PI;
    }
    
      // 应用运动控制
      if (strafe_mode == 0) {
          // 正常模式：前进/后退/转向
          for (int i = 0; i < 4; i++) {
              if (steers[i]) {
                  wb_motor_set_position(steers[i], steer_angle);
              }
          }
          
          // 设置驱动轮速度（注意左右侧轮子方向相反）
          if (drives[0]) wb_motor_set_velocity(drives[0], -speed);   // 左前
          if (drives[1]) wb_motor_set_velocity(drives[1], -speed);  // 右前
          if (drives[2]) wb_motor_set_velocity(drives[2], -speed);   // 左后
          if (drives[3]) wb_motor_set_velocity(drives[3], -speed);  // 右后
      }
      else if (strafe_mode == 1) {
          // 左移：所有轮子转向90度
          double strafe_angle = M_PI/2; // 90度
          for (int i = 0; i < 4; i++) {
              if (steers[i]) wb_motor_set_position(steers[i], strafe_angle);
          }
          
          // 设置驱动轮速度（对角线轮子同向）
          if (drives[0]) wb_motor_set_velocity(drives[0], -speed);
          if (drives[1]) wb_motor_set_velocity(drives[1], -speed);
          if (drives[2]) wb_motor_set_velocity(drives[2], -speed);
          if (drives[3]) wb_motor_set_velocity(drives[3], -speed);
      }
      else if (strafe_mode == 2) {
          // 右移：所有轮子转向-90度
          double strafe_angle = -M_PI/2; // -90度
          for (int i = 0; i < 4; i++) {
              if (steers[i]) wb_motor_set_position(steers[i], strafe_angle);
          }
          
          // 设置驱动轮速度（对角线轮子同向）
          if (drives[0]) wb_motor_set_velocity(drives[0], -speed);
          if (drives[1]) wb_motor_set_velocity(drives[1], -speed);
          if (drives[2]) wb_motor_set_velocity(drives[2], -speed);
          if (drives[3]) wb_motor_set_velocity(drives[3], -speed);
      }
      else if(strafe_mode == 3){
      double strafe_angle = -M_PI/4; // -45度
              
              
              if (steers[0]) wb_motor_set_position(steers[0], strafe_angle);
              if (steers[1]) wb_motor_set_position(steers[1], -strafe_angle);
              if (steers[2]) wb_motor_set_position(steers[2], -strafe_angle);
              if (steers[3]) wb_motor_set_position(steers[3], strafe_angle);
          
      // 设置驱动轮速度（对角线轮子同向）
          if (drives[0]) wb_motor_set_velocity(drives[0], speed);
          if (drives[1]) wb_motor_set_velocity(drives[1], -speed);
          if (drives[2]) wb_motor_set_velocity(drives[2], speed);
          if (drives[3]) wb_motor_set_velocity(drives[3], -speed);
      }
      
       else if(strafe_mode == 4){
      double strafe_angle = -M_PI/4; // -45度
              
              
              if (steers[0]) wb_motor_set_position(steers[0], strafe_angle);
              if (steers[1]) wb_motor_set_position(steers[1], -strafe_angle);
              if (steers[2]) wb_motor_set_position(steers[2], -strafe_angle);
              if (steers[3]) wb_motor_set_position(steers[3], strafe_angle);
          
      // 设置驱动轮速度（对角线轮子同向）
          if (drives[0]) wb_motor_set_velocity(drives[0], -speed);
          if (drives[1]) wb_motor_set_velocity(drives[1], speed);
          if (drives[2]) wb_motor_set_velocity(drives[2], -speed);
          if (drives[3]) wb_motor_set_velocity(drives[3], speed);
      }
      else if (strafe_mode == 5) {
        double vx_world = move_speed * cos(move_dir_angle);
        double vy_world = move_speed * sin(move_dir_angle);
        double vx, vy, w;
        if (rotate_mode) {
          // 开启自转模式，4轮全向轮状态（舵轮分别指向-45/45/-135/135度，全部朝外）
          // 修正右前和左后轮的驱动方向
          double omni_angles[4] = {-M_PI/4, M_PI/4, -3*M_PI/4, 3*M_PI/4}; // lf, rf, lb, rb
          int omni_signs[4] = {1, -1, -1, 1}; // lf, rf, lb, rb
          w = 5.0; // 只自转，不叠加平移
          double v_spin = MAX_SPEED; // 每个轮子自转速度设为最大速度
          for (int i = 0; i < 4; i++) {
            if (steers[i]) wb_motor_set_position(steers[i], omni_angles[i]);
            // 只自转分量，右前和左后轮速度取反
            double v = v_spin * omni_signs[i];
            if (drives[i]) wb_motor_set_velocity(drives[i], v);
          }
        } else {
          // 普通合成运动
          vx =  cos(robot_theta) * vx_world + sin(robot_theta) * vy_world;
          vy = -sin(robot_theta) * vx_world + cos(robot_theta) * vy_world;
          w = 0.0;
          double wheel_x[4] = {  0.2,  0.2, -0.2, -0.2 };
          double wheel_y[4] = {  0.2, -0.2,  0.2, -0.2 };
          double v_array[4];
          double v_max = 0.0;
          for (int i = 0; i < 4; i++) {
            double vix = vx * move_speed - w * wheel_y[i];
            double viy = vy * move_speed + w * wheel_x[i];
            double angle = atan2(viy, vix);
            double v = sqrt(vix*vix + viy*viy);
            // --- 最优舵轮转向 ---
            // 无传感器时，使用“上次下发的目标角度”模拟当前角度，避免突变
            static double last_steer_angle[4] = {0, 0, 0, 0};
            for (int i = 0; i < 4; i++) {
              double vix = vx * move_speed - w * wheel_y[i];
              double viy = vy * move_speed + w * wheel_x[i];
              double angle = atan2(viy, vix);
              double v = sqrt(vix*vix + viy*viy);
              double current_angle = last_steer_angle[i];
              // 归一化到[-pi, pi]
              angle = fmod(angle + M_PI, 2*M_PI); if (angle < 0) angle += 2*M_PI; angle -= M_PI;
              current_angle = fmod(current_angle + M_PI, 2*M_PI); if (current_angle < 0) current_angle += 2*M_PI; current_angle -= M_PI;
              double delta = angle - current_angle;
              // 归一化delta到[-pi, pi]
              delta = fmod(delta + M_PI, 2*M_PI); if (delta < 0) delta += 2*M_PI; delta -= M_PI;
              // 计算反向后的delta
              double delta_flip = delta > 0 ? delta - M_PI : delta + M_PI;
              // 只有反向后距离更短才允许反向，彻底消除180°突变
              if (fabs(delta_flip) < fabs(delta) - 1e-3) {
                angle += (delta > 0 ? -M_PI : M_PI);
                v = -v;
                // 再次归一化
                angle = fmod(angle + M_PI, 2*M_PI); if (angle < 0) angle += 2*M_PI; angle -= M_PI;
              }
              // 最终将angle归一化到[-pi/2, pi/2]
              if (angle > M_PI/2) angle -= M_PI;
              if (angle < -M_PI/2) angle += M_PI;
              last_steer_angle[i] = angle; // 记录本次下发的目标角度
              v_array[i] = v;
              if (fabs(v) > v_max) v_max = fabs(v);
              if (steers[i]) wb_motor_set_position(steers[i], angle);
            }
            // 归一化，保证每个轮子的速度不超过MAX_SPEED
            double scale = 1.0;
            if (v_max > MAX_SPEED) scale = MAX_SPEED / v_max;
            for (int i = 0; i < 4; i++) {
              if (drives[i]) wb_motor_set_velocity(drives[i], v_array[i] * scale);
            }
          }
        }
      }
      
      
      
  };
   
  /* Enter your cleanup code here */

  /* This is necessary to cleanup webots resources */
  wb_robot_cleanup();

  return 0;
}
